nanopyx.core.analysis.registration

  1from functools import cache
  2
  3import numpy as np
  4
  5from skimage.filters import window
  6
  7from .estimate_shift import GetMaxOptimizer
  8from .ccm_helper_functions import make_even_square
  9from .ccm import calculate_slice_ccm
 10from ..transform.interpolation_catmull_rom import Interpolator
 11
 12class Registration:
 13
 14    def __init__(self, image:np.ndarray, ref_image:np.ndarray):
 15        """
 16        Register an image against a reference image
 17        :param image: 2D array to register
 18        :param ref_image: 2D array to use as reference
 19        """
 20
 21        assert image.ndim == 2, "Image must be 2D"
 22        assert ref_image.ndim == 2, "Image must be 2D"
 23
 24        self.image = make_even_square(image[np.newaxis,:,:].astype(np.float32))[0,:,:]
 25        self.ref_image = make_even_square(ref_image[np.newaxis,:,:].astype(np.float32))[0,:,:]
 26
 27        self.original_dtype = image.dtype
 28
 29        self.w = self.image.shape[1]
 30        self.h = self.image.shape[0]
 31        self.wref = self.ref_image.shape[1] 
 32        self.href = self.ref_image.shape[0]
 33
 34        self.reg_result = {}
 35    
 36    
 37    def translation(self):
 38        """
 39        Registers the images considering only translation
 40        :return: registered image
 41        """
 42
 43        shifts, max_sim = self.phase_correlation(self.ref_image, self.image)
 44
 45        # The size of the CCM array is the same as the image
 46        y_shift = (self.h/2.0 - shifts[0])
 47        x_shift = (self.w/2.0 - shifts[1])
 48
 49        translated = Interpolator(self.image).shift(x_shift,y_shift).astype(self.original_dtype)
 50
 51        self.reg_result = {'Image':translated, 'Translation':(y_shift,x_shift), 'Scaling':None, 'Rotation':None, 'Max_Sim':max_sim}
 52
 53        return translated
 54    
 55
 56    def scaled_rotation(self):
 57        """
 58        Registers the images considering only rotation and isotropic scaling 
 59        """
 60        
 61        lpolar_image = Interpolator(self.image).polar(scale='log')
 62        lpolar_ref_image = Interpolator(self.ref_image).polar(scale='log')
 63        shifts, max_sim = self.phase_correlation(lpolar_ref_image, lpolar_image)
 64
 65        # Size of the polar transform is always (360,maxradius)
 66        h = 360
 67        w = np.hypot(self.w/2, self.h/2)
 68
 69        angle = -np.deg2rad((h/2 - shifts[0]))
 70        log_translation = (w/2-shifts[1]) * np.log(w) / w
 71        scale = np.exp(log_translation)
 72
 73        scaled = Interpolator(self.image).scale_xy(scale, scale)
 74        rotated = Interpolator(scaled).rotate(angle).astype(self.original_dtype)
 75
 76        self.reg_result = {'Image':rotated, 'Translation':None, 'Scaling':scaled, 'Rotation':angle, 'Max_Sim':max_sim}
 77
 78        return rotated
 79
 80    def scaling_rotation_translation(self):
 81        """
 82        Registers the images considering rotation, isotropic scaling and translation
 83        Based upon:
 84                An FFT-Based Technique for Translation,Rotation, 
 85                and Scale-Invariant Image Registration 
 86                B. Srinivasa Reddy and B. N. Chatterji
 87        """
 88        # Step 0: Prepare some heavily used vars
 89        h = 360
 90        w = np.hypot(self.w/2, self.h/2)
 91        highpass_filter = self.highpass_filter((self.h,self.w))
 92
 93        # Step 1: Prep the reference image for iteration
 94        windowed_ref_image = self.ref_image * window('hann', self.ref_image.shape)
 95        freq_ref_image = np.abs(np.fft.fftshift(np.fft.fft2(windowed_ref_image)) * highpass_filter).astype(np.float32)
 96        lpolar_ref_image = Interpolator(freq_ref_image).polar('log')
 97
 98        # Step 2: Iterate to find scale and angle
 99        total_angle = 0
100        total_scale = 1
101        iter_image = self.image
102
103        for iter in range(10): 
104            windowed_image = iter_image * window('hann', iter_image.shape)
105            f_image = np.abs(np.fft.fftshift(np.fft.fft2(windowed_image))*highpass_filter).astype(np.float32)
106            lpolar_image = Interpolator(f_image).polar('log')
107
108            shifts, max_sim_1 = self.phase_correlation(lpolar_ref_image, lpolar_image)
109
110            angle = -np.deg2rad((h/2 - shifts[0])) 
111            log_translation = (w/2-shifts[1]) * np.log(w) / w
112            scale = np.exp(-1*log_translation) # NEGATIVE SIGN HERE IS MANDATORY
113
114            total_angle += angle
115            total_scale *= scale
116
117            scaled = Interpolator(self.image).scale_xy(total_scale, total_scale)
118            rotated = Interpolator(scaled).rotate(total_angle)
119            iter_image = rotated
120     
121        # Step 3: Find translation AND the correct angle
122        # As outlined in the ref, in the frequency space 180-angle or angle have the same effect so there is ambiguity
123        # We test for both options and choose the one with the highest peak in the ccm
124
125        final_scaled_img = Interpolator(self.image).scale_xy(total_scale, total_scale)
126
127        final_option_1 = Interpolator(final_scaled_img).rotate(total_angle)
128        final_option_2 = Interpolator(final_scaled_img).rotate(np.pi+total_angle)
129
130        shifts_1, sim_1 = self.phase_correlation(self.ref_image, final_option_1)
131        shifts_2, sim_2 = self.phase_correlation(self.ref_image, final_option_2)
132
133        if sim_1 > sim_2:
134            final_iter_image = final_option_1
135            shifts = shifts_1
136            max_sim_2 = sim_1
137        else:
138            final_iter_image = final_option_2
139            shifts = shifts_2
140            max_sim_2 = sim_2
141            
142        y_shift = (self.h/2.0 - shifts[0])
143        x_shift = (self.w/2.0 - shifts[1])
144
145        translated = Interpolator(final_iter_image).shift(x_shift,y_shift).astype(self.original_dtype)
146
147        self.reg_result = {'Image':translated, 'Translation':(y_shift,x_shift), 'Scaling':total_scale, 'Rotation':total_angle, 'Max_Sim':(max_sim_1,max_sim_2)}
148
149        return translated
150
151    @staticmethod
152    def phase_correlation(im1:np.ndarray, im2:np.ndarray)->tuple:
153        """
154        Perform phase correlation between two images and return the shift that maximizes the overlap between the images
155        :param im1: 2D array of np.float32
156        :param im2: 2D array of np.float32
157        :return: coordinate tuple of the maximum point of the ccm and max value of the ccm
158        """
159
160        ccm = calculate_slice_ccm(im1, im2)
161        optimizer = GetMaxOptimizer(ccm)
162        shifts = optimizer.get_max()
163        maxsim = -optimizer.get_interpolated_px_value(shifts)
164
165        return shifts, maxsim
166
167    @staticmethod
168    @cache
169    def highpass_filter(shape)->np.ndarray:
170
171        n_row = shape[0]
172        n_col = shape[1]
173        row_freq_arr = np.fft.fftshift(np.fft.fftfreq(n_row))
174        col_freq_arr = np.fft.fftshift(np.fft.fftfreq(n_col))
175        row_f,col_f = np.meshgrid(row_freq_arr, col_freq_arr, indexing='ij')
176        X = np.cos(np.pi*row_f) * np.cos(np.pi*col_f) 
177        H = (1-X)*(2-X)
178
179        return H.astype(np.float32)